`include "defines.v"

module mem_stage(
    input wire              rst,

    input wire [63: 0]      inst_addr_i,
    input wire [31: 0]      inst_i,

    input wire              rd_w_ena_i,
    input wire [4 : 0]      rd_w_addr_i,
    input wire [`REG_BUS]   rd_w_data_i,

    input wire              csr_w_ena_i,
    input wire [11: 0]      csr_w_addr_i,
    input wire [`REG_BUS]   csr_w_data_i,
    
    input wire [`REG_BUS]   mem_addr_i,
    input wire [`REG_BUS]   mem_w_data_i,

    input wire [`REG_BUS]   mem_r_data_i,
    
    input wire                  mem_load_en_i,
    input wire                  mem_store_en_i,
    input wire[`MEM_SEL_BUS]    mem_sel_i,

    input wire[`RD_SEL_BUS]     rd_sel_i,

    input wire [`REG_BUS]   exception_type_i,
    input wire              valid_i,

    output wire [63: 0]     inst_addr_o,
    output wire [31: 0]     inst_o,

    output wire             rd_w_ena_o,
    output wire [4 : 0]     rd_w_addr_o,
    output wire [`REG_BUS]  rd_w_data_o,

    output wire             csr_w_ena_o,
    output wire [11: 0]     csr_w_addr_o,
    output wire [`REG_BUS]  csr_w_data_o,

    output wire [ 7: 0]     mem_w_mask_o,
    output wire [`REG_BUS]  mem_w_data_o,
    output wire             mem_w_en_o,
    output wire [`REG_BUS]  mem_addr_o,
    output wire             mem_ce_o,

    output wire [`REG_BUS]  exception_type_o,
    output wire             valid_o
);
    assign inst_addr_o = inst_addr_i;
    assign inst_o      = inst_i;

    assign rd_w_ena_o  = (rst == 1'h0) ? rd_w_ena_i : 1'b0;
    assign rd_w_addr_o = rd_w_addr_i;
    
    assign csr_w_ena_o = csr_w_ena_i;
    assign csr_w_addr_o= csr_w_addr_i;
    assign csr_w_data_o= csr_w_data_i;
    
    assign exception_type_o = exception_type_i;
    assign valid_o = valid_i;

    //exception
    wire freeze_flag;
    assign freeze_flag = exception_type_o == `ZERO_WORD ? 1'b0 : 1'b1;

    wire inst_lb_flag;
    wire inst_lh_flag;
    wire inst_lw_flag;
    wire inst_ld_flag;
    wire inst_lbu_flag;
    wire inst_lhu_flag;
    wire inst_lwu_flag;
    wire inst_sb_flag;
    wire inst_sh_flag;
    wire inst_sw_flag;
    wire inst_sd_flag;
    
    assign inst_lb_flag     = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LB ) ? 1'h1 : 1'h0;
    assign inst_lh_flag     = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LH ) ? 1'h1 : 1'h0;
    assign inst_lw_flag     = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LW ) ? 1'h1 : 1'h0;
    assign inst_ld_flag     = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LD ) ? 1'h1 : 1'h0;
    assign inst_lbu_flag    = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LBU) ? 1'h1 : 1'h0;
    assign inst_lhu_flag    = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LHU) ? 1'h1 : 1'h0;
    assign inst_lwu_flag    = (mem_load_en_i  == 1'b1) && (mem_sel_i == `MEM_LWU) ? 1'h1 : 1'h0;

    assign inst_sb_flag     = (mem_store_en_i == 1'b1) && (mem_sel_i == `MEM_SB ) ? 1'h1 : 1'h0;
    assign inst_sh_flag     = (mem_store_en_i == 1'b1) && (mem_sel_i == `MEM_SH ) ? 1'h1 : 1'h0;
    assign inst_sw_flag     = (mem_store_en_i == 1'b1) && (mem_sel_i == `MEM_SW ) ? 1'h1 : 1'h0;
    assign inst_sd_flag     = (mem_store_en_i == 1'b1) && (mem_sel_i == `MEM_SD ) ? 1'h1 : 1'h0;

    /****************************
    *load
    *****************************/
    wire[`REG_BUS] load_lb_result;
    wire[`REG_BUS] load_lbu_result;
    wire[`REG_BUS] load_lh_result;
    wire[`REG_BUS] load_lhu_result;
    wire[`REG_BUS] load_lw_result;
    wire[`REG_BUS] load_lwu_result;
    wire[`REG_BUS] load_ld_result;
    
    wire[`REG_BUS] load_result;

    
    assign load_lb_result  = {64{inst_lb_flag &  mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0]}} & {{56{mem_r_data_i[63]}}, mem_r_data_i[63:56]}
                           | {64{inst_lb_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{mem_r_data_i[55]}}, mem_r_data_i[55:48]}
                           | {64{inst_lb_flag &  mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0]}} & {{56{mem_r_data_i[47]}}, mem_r_data_i[47:40]}
                           | {64{inst_lb_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{mem_r_data_i[39]}}, mem_r_data_i[39:32]}
                           | {64{inst_lb_flag & ~mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0]}} & {{56{mem_r_data_i[31]}}, mem_r_data_i[31:24]}
                           | {64{inst_lb_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{mem_r_data_i[23]}}, mem_r_data_i[23:16]}
                           | {64{inst_lb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0]}} & {{56{mem_r_data_i[15]}}, mem_r_data_i[15: 8]}
                           | {64{inst_lb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{mem_r_data_i[ 7]}}, mem_r_data_i[ 7: 0]};

    assign load_lbu_result = {64{inst_lbu_flag &  mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[63:56]}
                           | {64{inst_lbu_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[55:48]}
                           | {64{inst_lbu_flag &  mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[47:40]}
                           | {64{inst_lbu_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[39:32]}
                           | {64{inst_lbu_flag & ~mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[31:24]}
                           | {64{inst_lbu_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[23:16]}
                           | {64{inst_lbu_flag & ~mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[15: 8]}
                           | {64{inst_lbu_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{56{1'b0}}, mem_r_data_i[ 7: 0]};

    assign load_lh_result  = {64{inst_lh_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{mem_r_data_i[63]}}, mem_r_data_i[63:48]}
                           | {64{inst_lh_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{mem_r_data_i[47]}}, mem_r_data_i[47:32]}
                           | {64{inst_lh_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{mem_r_data_i[31]}}, mem_r_data_i[31:16]}
                           | {64{inst_lh_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{mem_r_data_i[15]}}, mem_r_data_i[15: 0]};
    
    assign load_lhu_result = {64{inst_lhu_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{1'b0}}, mem_r_data_i[63:48]}
                           | {64{inst_lhu_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{1'b0}}, mem_r_data_i[47:32]}
                           | {64{inst_lhu_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{1'b0}}, mem_r_data_i[31:16]}
                           | {64{inst_lhu_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{48{1'b0}}, mem_r_data_i[15: 0]};

    assign load_lw_result  = {64{inst_lw_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{32{mem_r_data_i[63]}}, mem_r_data_i[63: 32]}
                           | {64{inst_lw_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{32{mem_r_data_i[31]}}, mem_r_data_i[31: 0]};

    assign load_lwu_result = {64{inst_lwu_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{32{1'b0}}, mem_r_data_i[63: 32]}
                           | {64{inst_lwu_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & {{32{1'b0}}, mem_r_data_i[31: 0]};

    assign load_ld_result  = {64{inst_ld_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0]}} & mem_r_data_i;

    assign load_result = load_lb_result | load_lbu_result | load_lh_result | load_lhu_result | load_lw_result | load_lwu_result | load_ld_result;
    
    assign rd_w_data_o = (mem_load_en_i == 1'b1) ? load_result : rd_w_data_i;
    /****************************
    *store
    *****************************/
    wire[ 7: 0]   store_sb_mask;
    wire[ 7: 0]   store_sh_mask;
    wire[ 7: 0]   store_sw_mask;
    wire[ 7: 0]   store_sd_mask;

    wire[`REG_BUS]   store_sb_data;
    wire[`REG_BUS]   store_sh_data;
    wire[`REG_BUS]   store_sw_data;
    wire[`REG_BUS]   store_sd_data;
    assign store_sb_mask = {8{inst_sb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0000_0001
                         | {8{inst_sb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0] }} & 8'b0000_0010
                         | {8{inst_sb_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0000_0100
                         | {8{inst_sb_flag & ~mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0] }} & 8'b0000_1000
                         | {8{inst_sb_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0001_0000
                         | {8{inst_sb_flag &  mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0] }} & 8'b0010_0000
                         | {8{inst_sb_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0100_0000
                         | {8{inst_sb_flag &  mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0] }} & 8'b1000_0000;

    assign store_sh_mask = {8{inst_sh_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0000_0011
                         | {8{inst_sh_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0000_1100
                         | {8{inst_sh_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0011_0000
                         | {8{inst_sh_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b1100_0000;

    assign store_sw_mask = {8{inst_sw_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b0000_1111
                         | {8{inst_sw_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b1111_0000;

    assign store_sd_mask = {8{inst_sd_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & 8'b1111_1111;

    assign store_sb_data = {64{inst_sb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & {56'h0, mem_w_data_i[7:0]}
                         | {64{inst_sb_flag & ~mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0] }} & {48'h0, mem_w_data_i[7:0],  8'h0}
                         | {64{inst_sb_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & {40'h0, mem_w_data_i[7:0], 16'h0}
                         | {64{inst_sb_flag & ~mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0] }} & {32'h0, mem_w_data_i[7:0], 24'h0}
                         | {64{inst_sb_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & {24'h0, mem_w_data_i[7:0], 32'h0}
                         | {64{inst_sb_flag &  mem_addr_i[2] & ~mem_addr_i[1] &  mem_addr_i[0] }} & {16'h0, mem_w_data_i[7:0], 40'h0}
                         | {64{inst_sb_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & { 8'h0, mem_w_data_i[7:0], 48'h0}
                         | {64{inst_sb_flag &  mem_addr_i[2] &  mem_addr_i[1] &  mem_addr_i[0] }} & {       mem_w_data_i[7:0], 56'h0};

    assign store_sh_data = {64{inst_sh_flag & ~mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & {48'h0, mem_w_data_i[15: 0]}
                         | {64{inst_sh_flag & ~mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & {32'h0, mem_w_data_i[15: 0], 16'h0}
                         | {64{inst_sh_flag &  mem_addr_i[2] & ~mem_addr_i[1] & ~mem_addr_i[0] }} & {16'h0, mem_w_data_i[15: 0], 32'h0}
                         | {64{inst_sh_flag &  mem_addr_i[2] &  mem_addr_i[1] & ~mem_addr_i[0] }} & {       mem_w_data_i[15: 0], 48'h0};

    assign store_sw_data = {64{inst_sw_flag & ~mem_addr_i[2] &  ~mem_addr_i[1] & ~mem_addr_i[0] }} & {32'h0, mem_w_data_i[31:0]}
                         | {64{inst_sw_flag &  mem_addr_i[2] &  ~mem_addr_i[1] & ~mem_addr_i[0] }} & {mem_w_data_i[31:0], 32'h0};

    assign store_sd_data = {64{inst_sd_flag & ~mem_addr_i[2] &  ~mem_addr_i[1] & ~mem_addr_i[0] }} & mem_w_data_i;
    ////////////////////////
    assign mem_w_mask_o = store_sb_mask
                        | store_sh_mask
                        | store_sw_mask
                        | store_sd_mask;

    assign mem_w_data_o = store_sb_data
                        | store_sh_data
                        | store_sw_data
                        | store_sd_data;

    // assign mem_w_en_o   = inst_sb_flag
    //                     | inst_sh_flag
    //                     | inst_sw_flag
    //                     | inst_sd_flag;

    assign mem_w_en_o   = mem_store_en_i;

    assign mem_addr_o   = mem_addr_i;

    assign mem_ce_o     = freeze_flag ? 1'b0 : inst_lb_flag
                                             | inst_lh_flag
                                             | inst_lw_flag
                                             | inst_ld_flag
                                             | inst_lbu_flag
                                             | inst_lhu_flag
                                             | inst_lwu_flag
                                             | inst_sb_flag
                                             | inst_sh_flag
                                             | inst_sw_flag
                                             | inst_sd_flag;
    
endmodule
